#  -*- coding: utf-8 -*-

from pymongo import ASCENDING,DESCENDING
from util.database import DB_CONN
from datetime import datetime, timedelta
import multiprocessing

def multi_computer(func,codes,args=()):
    process_num = 16
    code_multi_lst = [[] for i in range(process_num)]
    apply_results = []
    results = []
    for i in range(len(codes)):
        n = i % process_num
        code_multi_lst[n].append(codes[i])

    pool = multiprocessing.Pool(processes=process_num)
    for i in range(process_num):
        func_args = (code_multi_lst[i],) + args
        apply_results.append(pool.apply_async(func,args=func_args))
    pool.close() # 关闭进程池，表示不能再往进程池中添加进程，需要在join之前调用
    pool.join()
    for res in apply_results:
        results.append(res.get())

    print("The number of CPU is: %s " % str(multiprocessing.cpu_count()))
    return results

def get_trading_dates(begin_date=None, end_date=None):
    """
    获取指定日期范围的按照正序排列的交易日列表
    如果没有指定日期范围，则获取A股所有交易日

    :param begin_date: 开始日期
    :param end_date: 结束日期
    :return: 日期列表
    """
    # 开始日期，默认A股第一个交易日
    now = datetime.now()
    if begin_date is None:
        begin_date = "1990-12-19"

    # 结束日期默认为今天
    if end_date is None:
        end_date = now.strftime('%Y-%m-%d')
    daily_cursor = DB_CONN.daily.find(
        {'code': '999999', 'date': {'$gte': begin_date, '$lte': end_date}, 'index': True},
        sort=[('date', ASCENDING)],
        projection={'date': True, '_id': False})
    dates = [x['date'] for x in daily_cursor]

    return dates

def get_diff_dates(code=None,begin_date=None, end_date=None):
    """
    获取指定两个日期之间相差的交易日个数,如果不传入code代码，返回正常交易日差值，传入代码，则去除停牌日
    :param begin_date: 开始日期
    :param end_date: 结束日期
    :return: 交易日差个数
    """

    is_index = False
    if code == None:
        code = "999999"
        is_index = True

    daily_cursor = DB_CONN.daily.find(
        {'code': code, 'date': {'$gte': begin_date, '$lte': end_date}, 'index': is_index, "is_trading":True},
        sort=[('date', ASCENDING)],
        projection={'date': True, '_id': False})
    dates = [x['date'] for x in daily_cursor]
    return len(dates)


def get_code_trading_dates(code=None,is_index=False,begin_date=None, end_date=None):
    """
    获取指定两个日期之间相差的交易日个数,如果不传入code代码，返回正常交易日差值，传入代码，则去除停牌日
    :param begin_date: 开始日期
    :param end_date: 结束日期
    :return: 交易日差个数
    """

    if code == None:
        code = "999999"
        is_index = True

    daily_cursor = DB_CONN.daily.find(
        {'code': code, 'date': {'$gte': begin_date, '$lte': end_date}, 'index': is_index, "is_trading":True},
        sort=[('date', ASCENDING)],
        projection={'date': True, '_id': False})
    dates = [x['date'] for x in daily_cursor]
    return dates

def judge_code_trading_date(code=None,is_index=False,date=None):
    if code == None:
        code = "999999"
        is_index = True
    #print(code,date)
    daily_cursor = DB_CONN.daily.find(
        {'code': code, 'date': {'$lte': date}, 'index': is_index, "is_trading":True},
        sort=[('date', ASCENDING)],
        projection={'date': True, '_id': False})
    dates = [x['date'] for x in daily_cursor]
    #print(dates)
    return dates[-1]

def calc_negative_diff_dates(code=None,is_index=False,date=None, delta_days=0):
    """
    计算指定日期相差指定交易日得到的交易日期
    :param date: 日期
    :param delta_days: 交易日差个数,负的向前，正的向后
    :return: 差值交易日期
    """
    delta_date_index = 0
    if code == None:
        code = "999999"
        is_index = True

    daily_cursor = DB_CONN.daily.find(
        {'code': code, 'index': is_index, "is_trading":True,'date': {'$lte': date}},
        sort=[('date', ASCENDING)],
        projection={'date': True, '_id': False})
    dates = [x['date'] for x in daily_cursor]
    print(dates)
    if date in dates:
        delta_date_index = dates.index(date) + delta_days
    else:
        #Todo:这里应该找到离date最近的前一个日期，然后去计算偏移，而不是直接用序列最后一个来计算
        delta_date_index = dates.index(dates[-1]) + delta_days

    if delta_date_index >= 0 :
        if len(dates) > delta_date_index :
            delta_date = dates[delta_date_index]
        else:
            delta_date = dates[-1]
    else:
        delta_date = dates[0]

    return delta_date

def calc_positive_diff_dates(code=None,is_index=False,date=None, delta_days=0):
    """
    计算指定日期相差指定交易日得到的交易日期
    :param date: 日期
    :param delta_days: 交易日差个数,负的向前，正的向后
    :return: 差值交易日期
    """
    delta_date_index = 0
    if code == None:
        code = "999999"
        is_index = True

    daily_cursor = DB_CONN.daily.find(
        {'code': code, 'index': is_index, "is_trading":True,'date': {'$gte': date}},
        sort=[('date', ASCENDING)],
        projection={'date': True, '_id': False})
    dates = [x['date'] for x in daily_cursor]

    if date in dates:
        delta_date_index = dates.index(date) + delta_days
    else:
        #Todo:这里应该找到离date最近的前一个日期，然后去计算偏移，而不是直接用序列最后一个来计算
        delta_date_index = dates.index(dates[-1]) + delta_days

    if delta_date_index >= 0 :
        if len(dates) > delta_date_index :
            delta_date = dates[delta_date_index]
        else:
            delta_date = dates[-1]
    else:
        delta_date = dates[0]

    return delta_date

def get_all_codes_date(date=None):
    """
    获取某个交易日的所有股票代码列表，如果没有指定日期，则从当前日期一直向前找，直到找到有
    数据的一天，返回的即是那个交易日的股票代码列表

    :param date: 日期
    :return: 股票代码列表
    """

    #todo 没考虑退市情况，算法不对
    datetime_obj = datetime.now()
    if date is None:
        date = datetime_obj.strftime('%Y-%m-%d')

    codes = []
    while len(codes) == 0:
        code_cursor = DB_CONN.stable.find(
            {'date': date},
            projection={'code': True, '_id': False})

        codes = [x['code'] for x in code_cursor]

        datetime_obj = datetime_obj - timedelta(days=1)
        date = datetime_obj.strftime('%Y-%m-%d')

    return codes

def get_all_codes_trading_date(date=None):
    """
    获取某个交易日的所有股票代码列表，如果没有指定日期，则从当前日期一直向前找，直到找到有
    数据的一天，返回的即是那个交易日的股票代码列表

    :param date: 日期
    :return: 股票代码列表
    """

    datetime_obj = datetime.now()
    if date is None:
        date = datetime_obj.strftime('%Y-%m-%d')

    codes = []
    DB_CONN.daily_qfq.create_index([('date', 1), ('index', 1),('is_trading',1)])
    code_cursor = DB_CONN.daily_qfq.find(
        {'date': date, 'index': False, "is_trading": True},
        projection={'code': True, '_id': False})

    codes = [x['code'] for x in code_cursor]

    return codes

def get_all_trading_codes_data_trading_date(date=None):
    """
    获取某个交易日的所有有交易的股票个数和成交量之和，如果没有指定日期，则从当前日期一直向前找，直到找到有
    数据的一天，返回的即是那个交易日的股票代码列表

    :param date: 日期
    :return: 股票代码列表
    """

    datetime_obj = datetime.now()
    if date is None:
        date = datetime_obj.strftime('%Y-%m-%d')

    DB_CONN.daily_qfq.create_index([('date', 1), ('index', 1),('is_trading',1)])
    code_cursor = DB_CONN.daily_qfq.find(
        {'date': date, 'index': False, "is_trading": True},
        projection={'code': True, '_id': False})

    codes_num = len(list(code_cursor))
    total_amount= 0
    #上证数据
    sz_data =  DB_CONN.daily.find(
        {'code':"999999",'date': date, 'index': True, "is_trading": True},
        projection={'_id': False})
    #print(sz_data[0]['amount'])
    total_amount += sz_data[0]['amount']
    #深成数据
    sc_data = DB_CONN.daily.find(
        {'code':"399001",'date': date, 'index': True, "is_trading": True},
        projection={'_id': False})
    total_amount += sc_data[0]['amount']
    total_amount = round(total_amount / 100000000, 2)
    return codes_num,total_amount

def get_sub_industry(code):
    sub_industry = ""
    try:
        stable_cursor = DB_CONN['stable'].find(
            {'code': code, 'index': False},
            projection={'_id': False},
            batch_size=1000)

        sub_industry = stable_cursor[0]['sub_industry']
    except:
        print(f"get_sub_industry error code:{code}")

    return sub_industry

def get_choice_block_mom_value(block_name,date):
    mom_value = 0
    momentum_cursor = DB_CONN['momentum'].find(
        {'name':block_name,'date': date, 'type': 'sub_industry', "origin": 'choice'},
        projection={'code_list': False, '_id': False})
    if momentum_cursor.count() > 0:
        mom_value = momentum_cursor[0]['momentum_value']

    return mom_value

def get_choice_block_mom_rank(block_name,date):
    mom_rank = 0
    momentum_cursor = DB_CONN['momentum'].find(
        {'name':block_name,'date': date, 'type': 'sub_industry', "origin": 'choice'},
        projection={'code_list': False, '_id': False})
    if momentum_cursor.count() > 0:
        mom_rank = momentum_cursor[0]['rank']

    return mom_rank

def get_all_codes():
    """
    获取当前数据库里的全部股票代码
    :return: 股票代码列表
    """
    code_cursor = DB_CONN.stable.find(
        {"index":False},
        projection={"code":True,'_id': False})

    codes = [x['code'] for x in code_cursor]
    #codes.remove('510300')
    #codes.remove('510500')
    #codes.remove('510310')
    #codes.remove('510900')
    #codes.remove('512400')
    #codes.remove('512880')
    #codes.remove('159915')
    return codes

def get_code_name(code,index):
    try:
        code_cursor = DB_CONN.stable.find(
            {"code":code,"index":index},
            projection={ '_id': False})
        name =code_cursor[0]['name']
    except:
        return "NotExists"
    return name

def get_all_indexes_date(date=None):
    """
    获取当前数据库里的全部指数代码
    :return: 股票代码列表
    """
    datetime_obj = datetime.now()
    if date is None:
        date = datetime_obj.strftime('%Y-%m-%d')

    codes = []
    while len(codes) == 0:
        code_cursor = DB_CONN.daily.find(
            {"index":True,'date': date},
            projection={'code': True, '_id': False})

        codes = [x['code'] for x in code_cursor]

        datetime_obj = datetime_obj - timedelta(days=1)
        date = datetime_obj.strftime('%Y-%m-%d')

    return codes

def get_all_indexes():
    """
    获取当前数据库里的全部股票代码
    :return: 股票代码列表
    """
    code_cursor = DB_CONN.stable.find(
        {"index":True},
        projection={"code":True,'_id': False})

    codes = [x['code'] for x in code_cursor]

    return codes

def get_index_cons_code_list(index_code,date):
    trading_date = judge_code_trading_date(date = date)
    index_cursor = DB_CONN.index_cons.find({'index_code': index_code,'update_date': {'$gte': trading_date}},
                                           sort=[('update_date', DESCENDING)],
                                           projection={'_id': False})
    code_list = list()
    for code_info in index_cursor[0]['code_list']:
        code_list.append(code_info['code'])
    return code_list


if __name__ == '__main__':
    #get_all_codes()
    #get_index_cons_code_list('000300',"2021-01-09")
    get_all_trading_codes_data_trading_date("2021-03-09")